'''
根据调和常数计算以月或者日为时间分辨率的平均潮位
'''

import sys
import calendar
import math
import os
import numpy as np
import pandas as pd
import time

root_path = r"__file__"                                                                             # 添加pydggs路径
if root_path not in sys.path:
    sys.path.insert(0, root_path)

from PyDGGS.PyDGGS import *
from common import NAN, ONE_HOUR_TO_SECOND, HOURS_PER_DAY,σ, tide_start_time, cal_f_u, cal_v        # 导入常量

time_resolution = 'daily'                                                                           # 设置时间分辨率:   月：monthly，日：daily

input_dir = os.path.join(os.path.dirname(r'__file__'))                                              # 数据位置
hg_output_path = os.path.join(input_dir, 'output', 'hg', 'hg.json')                                 # 读取提取出来的调和常数
daily_tide_elev_output_path = os.path.join(input_dir, 'output', 'daily_tide')                       # 存储反算出来的日平均潮高数据
monthly_tide_elev_output_path = os.path.join(input_dir, 'output', 'monthly_tide')                   # 存储反算出来的月平均潮高数据


def calElevation(waves_hg, wave_names, date_time):
    '''
    根据调和常数HG计算平均潮位

    :param waves_hg:    分潮的调和常数
    :param wave_names:  分潮名称
    :param date_time:   计算分潮的时间
    :return:            潮位
    '''
    start_hours = (date_time - tide_start_time).total_seconds() / ONE_HOUR_TO_SECOND  # 当天UTC 0时刻

    # 计算潮位，然后求取平均值
    day_result = np.zeros((HOURS_PER_DAY,), dtype=np.float32)
    for i in range(HOURS_PER_DAY):
        hours = start_hours + i
        wave_fu_dict = cal_f_u(hours)
        wave_v_dict = cal_v(hours)
        # 计算潮位
        result = []
        for wave in wave_names:
            wave_name = str(wave)
            wave_h = waves_hg[wave_name]['H']
            wave_g = waves_hg[wave_name]['G']

            wave_fcosu = wave_fu_dict[wave_name]['fcosu']
            wave_fsinu = wave_fu_dict[wave_name]['fsinu']

            wave_f = math.sqrt(wave_fcosu * wave_fcosu + wave_fsinu * wave_fsinu)
            angle = math.degrees(math.atan(wave_fsinu / wave_fcosu))
            # 根据波浪正弦余弦确定角度
            if wave_fsinu >= 0 >= wave_fcosu:
                wave_u = 180 - abs(angle)
            elif wave_fsinu <= 0 and wave_fcosu <= 0:
                wave_u = abs(angle) + 180
            elif wave_fsinu <= 0 <= wave_fcosu:
                wave_u = 360 - abs(angle)
            else:
                wave_u = angle

            wave_v = wave_v_dict[wave_name]
            wave_elev = wave_f * wave_h * math.cos(math.radians(wave_v + wave_u - wave_g))

            result.append(wave_elev)

        result = np.array(result)
        day_result[i] = np.sum(result)

    return float(np.mean(day_result))


def calElev(hg_file, date_time, day):
    """
    获取指定日期的平均潮位（24小时平均）

    :param hg_file:          GridTIFF格式的调和常数文件路径
    :param date_time:        计算日期
    :param day:              是否计算单日（非空字符串表示计算单日）
    :return:                 平均潮位和潮位列表
    """
    wave_names = ['K2', 'N2', 'S2', 'M2', 'Q1', 'O1', 'K1', 'P1']
    
    # 分潮参数配置 - 按照波段顺序
    tidal_constituents = [
        ('K2_H', 'K2_G'),  # 波段0,1
        ('N2_H', 'N2_G'),  # 波段2,3
        ('S2_H', 'S2_G'),  # 波段4,5
        ('M2_H', 'M2_G'),  # 波段6,7
        ('Q1_H', 'Q1_G'),  # 波段8,9
        ('O1_H', 'O1_G'),  # 波段10,11
        ('K1_H', 'K1_G'),  # 波段12,13
        ('P1_H', 'P1_G')   # 波段14,15
    ]
    
    # 加载GridTIFF格式的调和常数
    hg_layer = DGGSFieldLayerDouble()
    if not hg_layer.loadFromFile(hg_file):
        print(f"错误：无法加载调和常数文件 {hg_file}")
        return None, []
    
    level = hg_layer.getLevel()
    grid_type = hg_layer.getGridType()
    element_type = hg_layer.getElementType()
    
    # 确定计算类型和输出信息
    if day != '':
        time_info = f"{date_time.year}年{date_time.month}月{date_time.day}日"
    else:
        time_info = f"{date_time.year}年{date_time.month}月"
    
    print(f"正在获取{time_info}的平均潮位")
    
    elev_list = []
    row_size = rowMaxSize(level)
    col_size = fColMaxSize(level, grid_type)
    
    # 遍历所有格元
    for base_partition in range(1, 11):
        print(f"处理基础菱形 {base_partition}/10")
        for row in range(row_size):
            for col in range(col_size):
                code = DGGSElementCode()
                code.elementType = element_type
                code.level = level
                code.basePartition = base_partition
                code.row = row
                code.col = col
                
                # 从16个波段读取调和常数
                waves_hg = {}
                valid_data = True
                
                for i, wave_name in enumerate(wave_names):
                    h_band = i * 2      # H值在偶数波段
                    g_band = i * 2 + 1  # G值在奇数波段
                    
                    h_val = hg_layer.getValue(code, h_band)
                    g_val = hg_layer.getValue(code, g_band)
                    
                    # 处理无效值
                    if h_val == noData or g_val == noData:
                        valid_data = False
                        break
                    else:
                        waves_hg[wave_name] = {'H': h_val, 'G': g_val}
                
                # 如果数据无效，跳过该格元
                if not valid_data:
                    continue
                
                # 计算潮位
                if day != '':
                    elev = calElevation(waves_hg, wave_names, date_time)
                else:
                    year, month = date_time.year, date_time.month
                    _, num_days = calendar.monthrange(year, month)
                    all_days = pd.date_range(f'{year}-{month:02d}-01', periods=num_days)
                    
                    elev_sum = 0
                    valid_days = 0
                    for date in all_days:
                        daily_elev = calElevation(waves_hg, wave_names, date)
                        if not math.isnan(daily_elev):
                            elev_sum += daily_elev
                            valid_days += 1
                    
                    if valid_days > 0:
                        elev = elev_sum / valid_days
                    else:
                        elev = NAN
                
                # 计算格元ID并添加到结果列表
                cell_id = code2Id(code, grid_type)
                elev_list.append((cell_id, elev))
    
    print(f"计算{time_info}的潮位完成，有效格元数: {len(elev_list)}")
    return level, elev_list


def saveElevFile(hg_file, year, month, day, elev_tiff_dir):
    """
    保存潮高

    :param hg_file:          GridTIFF格式的调和常数文件路径
    :param year:             年
    :param month:            月
    :param day:              日
    :param elev_tiff_dir:    平均潮位GridTIFF保存目录
    """
    data_type = DGGS_Float64
    element_type = DGGSElementType_Cell
    grid_type = DGGSGridType_DGGS_ISEA4H
    band_count = 1
    
    # 根据day参数确定计算方式和保存路径
    if day:
        current_time = pd.to_datetime(f'{year}-{month}-{day}')
        save_suffix = f"{year}{month}{day}"
    else:
        current_time = pd.to_datetime(f'{year}-{month}-15')
        save_suffix = f"{year}{month}"
    
    level, elevs = calElev(hg_file, current_time, day)
    
    if level is None:
        print("错误：无法计算潮位")
        return
    
    field_layer = DGGSFieldLayerDouble(
        band_count, data_type, element_type, grid_type, level
    )
    
    # 初始化所有网格单元为NAN
    row_size = rowMaxSize(level)
    col_size = fColMaxSize(level, grid_type)
    
    for row in range(row_size):
        for col in range(col_size):
            for base_partition in range(1, 11):
                for band_id in range(band_count):
                    code = DGGSElementCode()
                    code.level = level
                    code.basePartition = base_partition
                    code.elementType = element_type
                    code.row = row
                    code.col = col
                    field_layer.setValue(code, NAN, band_id)
    
    print("设置潮高值到格网单元...")
    
    # 设置潮高值到对应网格单元
    for cell_id, elev in elevs:
        code = DGGSElementCode()
        code.level = level
        code.elementType = DGGSElementType_Cell
        id2Code(cell_id, code, grid_type)
        field_layer.setValue(code, elev, 0)
    
    # 构建输出文件路径 - 修正为目录路径
    output_dir = os.path.join(elev_tiff_dir, save_suffix)
    os.makedirs(output_dir, exist_ok=True)
    
    field_layer.saveToFile(output_dir)
    print(f"潮高文件已保存到: {output_dir}")


def main():

    months = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]        # 月份

    if time_resolution == 'daily':    # 时间分辨率是日
        years = [2015, 2016]                                # 年份
        for year in years:
            for month in months:
                _, days_in_month = calendar.monthrange(year, month)
                for day in range(1, days_in_month + 1):
                    year_str = str(year)
                    month_str = f'{month:02d}'
                    day_str = f'{day:02d}'

                    os.makedirs(os.path.join(daily_tide_elev_output_path, f"{year_str}{month_str}{day_str}"))
                    saveElevFile(hg_output_path, year_str, month_str, day_str, daily_tide_elev_output_path)

    
    elif time_resolution == 'monthly':    # 时间分辨率是月
        years = [2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019]     # 年份
        day = ''
        for year in years:
            for month in months:
                year_str = str(year)
                month_str = f'{month:02d}'

                os.makedirs(os.path.join(monthly_tide_elev_output_path, f"{year_str}{month_str}"))
                saveElevFile(hg_output_path, year_str, month_str, day, monthly_tide_elev_output_path)

    else:
        print("未知的时间分辨率")

if __name__ == '__main__':
    start_time = time.time()
    main()
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"程序执行时间: {execution_time} 秒")
